import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from .config import *
from .data_loader import process_images
import datetime
import os

def evaluate_model(model, test_data, times_test):
    """评估模型性能"""
    images, pv_log, pv_pred = test_data
    if config.MODEL_SELECT == "CNN_LSTM":
        loss = model.evaluate([images, pv_log], pv_pred, verbose=0)
        predictions = np.squeeze(model.predict(([images, pv_log]), batch_size=200, verbose=1))
    else:
        loss = model.evaluate(pv_log, pv_pred, verbose=0)
        predictions = np.squeeze(model.predict((pv_log), batch_size=200, verbose=1))
    rmse = np.sqrt(np.mean((predictions - pv_pred)**2))
    
    dates = np.array([t.date() for t in times_test])
    
    sunny_mask = np.isin(dates, [datetime.date(*d) for d in config.SUNNY_DATES])
    
    metrics = {
        'overall': calculate_metrics(predictions, pv_pred),
        'sunny': calculate_metrics(predictions[sunny_mask], pv_pred[sunny_mask]),
        'cloudy': calculate_metrics(predictions[~sunny_mask], pv_pred[~sunny_mask])
    }
    
    return metrics, predictions

def calculate_metrics(pred, true):
    mse = np.mean((pred - true)**2)
    mae = np.mean(np.abs(pred - true))
    return {'RMSE': np.sqrt(mse), 'MAE': mae, 'MSE': mse}

def plot_results(times, true, pred, sunny_dates, cloudy_dates):
    """可视化预测结果"""
    # 转换日期格式并生成mask
    dates = np.array([t.date() for t in times])
    
    # 创建基准日期用于时间轴显示
    base_date = datetime.datetime(2000, 1, 1)
    hours = np.array([base_date.replace(
        hour=t.hour,
        minute=t.minute,
        second=t.second
    ) for t in times])
    
    # 转换为matplotlib的日期格式
    hours_num = mdates.date2num(hours)
    
    # 确保数据转为numpy数组
    true = np.array(true).flatten()
    pred = np.array(pred).flatten()
    
    # 晴天/阴天颜色设置
    colors = {
        'Sunny': '#FF6B6B',   # 珊瑚红
        'Cloudy': '#4ECDC4'   # 蓝绿色
    }
    
    # 计算需要绘制的总子图数量
    total_plots = len(sunny_dates) + len(cloudy_dates)
    fig, axs = plt.subplots(total_plots, 1, figsize=(12, 5 * total_plots), sharex=True)
    
    # 如果只有一个子图，确保 axs 是可迭代对象
    if total_plots == 1:
        axs = [axs]
    
    def draw_plot(ax, mask, title, color):
        # 绘制真实值和预测值
        ax.plot(hours_num[mask], true[mask],
                label='Ground Truth',
                color='#2c3e50',  # 深灰色
                alpha=0.8,
                linewidth=2)
        ax.plot(hours_num[mask], pred[mask],
                label='Prediction',
                color=color,
                linestyle='--',
                linewidth=2)
        
        # 设置时间格式
        ax.xaxis.set_major_locator(mdates.HourLocator(interval=4))  # 修改间隔为每4小时一个刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
        
        # 添加图例和标题
        ax.legend(loc='upper right')
        ax.set_title(title, pad=20)
        
        # 设置坐标轴标签
        ax.set_ylabel('PV Output (kW)')
        
        # 添加网格线
        ax.grid(True, alpha=0.3)
    
    # 绘制晴天的性能图
    idx = 0
    for date in sunny_dates:
        mask = dates == date
        draw_plot(axs[idx], mask, f'Sunny Day: {date.strftime("%Y-%m-%d")}', colors['Sunny'])
        idx += 1
    
    # 绘制阴天的性能图
    for date in cloudy_dates:
        mask = dates == date
        draw_plot(axs[idx], mask, f'Cloudy Day: {date.strftime("%Y-%m-%d")}', colors['Cloudy'])
        idx += 1
    
    # 最后一个子图添加X轴标签
    axs[-1].set_xlabel('Time of Day')
    
    # 调整布局
    plt.tight_layout(pad=3.0)
    
    output_path = os.path.join(config.output_folder, 'forecast_comparison.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved visualization to {output_path}")
    

def plot_sunny_results(times, true, pred, sunny_dates):
    """可视化晴天天气情况下的预测结果"""
    # 转换日期格式并生成mask
    dates = np.array([t.date() for t in times])
    
    # 创建基准日期用于时间轴显示
    base_date = datetime.datetime(2000, 1, 1)
    hours = np.array([base_date.replace(
        hour=t.hour,
        minute=t.minute,
        second=t.second
    ) for t in times])
    
    # 转换为matplotlib的日期格式
    hours_num = mdates.date2num(hours)
    
    # 确保数据转为numpy数组
    true = np.array(true).flatten()
    pred = np.array(pred).flatten()
    
    # 晴天颜色设置
    color = '#FF6B6B'   # 珊瑚红
    
    # 计算需要绘制的总子图数量
    total_plots = len(sunny_dates)
    fig, axs = plt.subplots(total_plots, 1, figsize=(12, 5 * total_plots), sharex=True)
    
    if total_plots == 1:
        axs = [axs]
    
    def draw_plot(ax, mask, title, color):
        # 绘制真实值和预测值
        ax.plot(hours_num[mask], true[mask],
                label='Ground Truth',
                color='#2c3e50',  # 深灰色
                alpha=0.8,
                linewidth=2)
        ax.plot(hours_num[mask], pred[mask],
                label='Prediction',
                color=color,
                linestyle='--',
                linewidth=2)
        
        # 设置时间格式
        ax.xaxis.set_major_locator(mdates.HourLocator(interval=4))  # 修改间隔为每4小时一个刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
        
        # 添加图例和标题
        ax.legend(loc='upper right')
        ax.set_title(title, pad=20)
        
        # 设置坐标轴标签
        ax.set_ylabel('PV Output (kW)')
        
        # 添加网格线
        ax.grid(True, alpha=0.3)
    
    # 绘制晴天的性能图
    for idx, date in enumerate(sunny_dates):
        mask = dates == date
        draw_plot(axs[idx], mask, f'Sunny Day: {date.strftime("%Y-%m-%d")}', color)
    
    # 最后一个子图添加X轴标签
    axs[-1].set_xlabel('Time of Day')
    
    # 调整布局
    plt.tight_layout(pad=3.0)
    
    output_path = os.path.join(config.output_folder, 'sunny_forecast_comparison.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved sunny day visualization to {output_path}")

def plot_cloudy_results(times, true, pred, cloudy_dates):
    """可视化阴天天气情况下的预测结果"""
    # 转换日期格式并生成mask
    dates = np.array([t.date() for t in times])
    
    # 创建基准日期用于时间轴显示
    base_date = datetime.datetime(2000, 1, 1)
    hours = np.array([base_date.replace(
        hour=t.hour,
        minute=t.minute,
        second=t.second
    ) for t in times])
    
    # 转换为matplotlib的日期格式
    hours_num = mdates.date2num(hours)
    
    # 确保数据转为numpy数组
    true = np.array(true).flatten()
    pred = np.array(pred).flatten()
    
    # 阴天颜色设置
    color = '#4ECDC4'   # 蓝绿色
    
    # 计算需要绘制的总子图数量
    total_plots = len(cloudy_dates)
    fig, axs = plt.subplots(total_plots, 1, figsize=(12, 5 * total_plots), sharex=True)
    
    if total_plots == 1:
        axs = [axs]
    
    def draw_plot(ax, mask, title, color):
        # 绘制真实值和预测值
        ax.plot(hours_num[mask], true[mask],
                label='Ground Truth',
                color='#2c3e50',  # 深灰色
                alpha=0.8,
                linewidth=2)
        ax.plot(hours_num[mask], pred[mask],
                label='Prediction',
                color=color,
                linestyle='--',
                linewidth=2)
        
        # 设置时间格式
        ax.xaxis.set_major_locator(mdates.HourLocator(interval=4))  # 修改间隔为每4小时一个刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M'))
        
        # 添加图例和标题
        ax.legend(loc='upper right')
        ax.set_title(title, pad=20)
        
        # 设置坐标轴标签
        ax.set_ylabel('PV Output (kW)')
        
        # 添加网格线
        ax.grid(True, alpha=0.3)
    
    # 绘制阴天的性能图
    for idx, date in enumerate(cloudy_dates):
        mask = dates == date
        draw_plot(axs[idx], mask, f'Cloudy Day: {date.strftime("%Y-%m-%d")}', color)
    
    # 最后一个子图添加X轴标签
    axs[-1].set_xlabel('Time of Day')
    
    # 调整布局
    plt.tight_layout(pad=3.0)
    
    output_path = os.path.join(config.output_folder, 'cloudy_forecast_comparison.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved cloudy day visualization to {output_path}")



